from typing import TypeVar

import numpy as np

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.processing import DelegatingDataset, SlicedDataset

AnySample = TypeVar('AnySample')


class SlicingShardedDataset(DelegatingDataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], rank: int, size: int):
        ranks_head = (len(dataset) / size) * np.arange(size)
        ranks_head = np.asarray(ranks_head, dtype=int)
        ranks_tail = np.concat([ranks_head[1:], [len(dataset)]])
        rank_head = int(ranks_head[rank])
        ranks_tail = int(ranks_tail[rank])
        super().__init__(SlicedDataset(dataset, slice(rank_head, ranks_tail, None)))
        self._dataset = dataset
        self._rank = rank
        self._size = size
        self._head = rank_head
        self._tail = ranks_tail

    @property
    def datasets(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def rank(self) -> int:
        return self._rank

    @property
    def size(self) -> int:
        return self._size

    @property
    def head(self) -> int:
        return self._head

    @property
    def tail(self) -> int:
        return self._tail
